{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notebook written by [Zhedong Zheng](https://github.com/zhedongzheng)\n",
    "\n",
    "![title](img/char_embed.gif)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "from collections import Counter\n",
    "\n",
    "import tensorflow as tf\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "VOCAB_SIZE = 5000\n",
    "MAX_LEN = 400\n",
    "BATCH_SIZE = 32\n",
    "EMBED_DIM = 50\n",
    "FILTERS = 250\n",
    "N_CLASS = 2\n",
    "N_EPOCH = 2\n",
    "LR = {'start': 5e-3, 'end': 5e-4, 'steps': 1500}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_idx2word(_index_from=3):\n",
    "    word2idx = tf.keras.datasets.imdb.get_word_index()\n",
    "    word2idx = {k:(v+_index_from) for k,v in word2idx.items()}\n",
    "    word2idx[\"<pad>\"] = 0\n",
    "    word2idx[\"<start>\"] = 1\n",
    "    word2idx[\"<unk>\"] = 2\n",
    "    idx2word = {idx: w for w, idx in word2idx.items()}\n",
    "    return idx2word, word2idx\n",
    "\n",
    "\n",
    "def make_feats(X, char2idx, word2idx, MAX_WORD_LEN):\n",
    "    char_t = np.zeros([len(X), MAX_LEN, MAX_WORD_LEN], dtype=np.int32)\n",
    "    word_t = np.zeros([len(X), MAX_LEN], dtype=np.int32)\n",
    "    for i, sent in tqdm(enumerate(X), total=len(X), ncols=70):\n",
    "        for j, w in enumerate(sent[-MAX_LEN:]):\n",
    "            if j < MAX_LEN:\n",
    "                word_t[i, j] = word2idx[''.join(w)]\n",
    "                for k, c in enumerate(w):\n",
    "                    char_t[i, j, k] = char2idx.get(c, char2idx['<unk>'])\n",
    "    return {'char': char_t, 'word': word_t}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def char_embedding(x, params, batch_size):\n",
    "    char2vec = tf.get_variable('char2vec', [params['char_vocab_size'], EMBED_DIM])\n",
    "    char2vec = tf.concat([tf.zeros([1, EMBED_DIM]), char2vec[1:, :]], axis=0)\n",
    "    x = tf.nn.embedding_lookup(char2vec, x)\n",
    "    x = tf.reshape(x, [batch_size*MAX_LEN, params['max_word_len'], EMBED_DIM])\n",
    "    x = tf.layers.conv1d(x, EMBED_DIM, kernel_size=5, activation=tf.nn.relu)\n",
    "    x = tf.reduce_max(x, 1)\n",
    "    x = tf.reshape(x, [batch_size, MAX_LEN, EMBED_DIM])\n",
    "    return x\n",
    "\n",
    "\n",
    "def word_embedding(x):\n",
    "    word2vec = tf.get_variable('word2vec', [VOCAB_SIZE, EMBED_DIM])\n",
    "    word2vec = tf.concat([tf.zeros([1, EMBED_DIM]), word2vec[1:, :]], axis=0)\n",
    "    return tf.nn.embedding_lookup(word2vec, x)\n",
    "\n",
    "\n",
    "def highway_layer(inputs):\n",
    "    size = inputs.get_shape()[-1].value\n",
    "    x1, x2 = tf.split(tf.layers.dense(inputs, 2*size), 2, -1)\n",
    "    flow = tf.nn.relu(x1)\n",
    "    gate = tf.sigmoid(x2)\n",
    "    return gate * flow + (1 - gate) * inputs\n",
    "\n",
    "\n",
    "def forward(feats, mode, params):\n",
    "    is_training = (mode == tf.estimator.ModeKeys.TRAIN)\n",
    "    batch_size = tf.shape(feats['char'])[0]\n",
    "    \n",
    "    x = tf.concat([word_embedding(feats['word']),\n",
    "                   char_embedding(feats['char'], params, batch_size),], axis=-1)\n",
    "    x = tf.layers.dropout(x, 0.2, training=is_training)\n",
    "    x = highway_layer(x)\n",
    "    \n",
    "    feat_map = []\n",
    "    for k_size in [3, 4, 5]:\n",
    "        _x = tf.layers.conv1d(x, FILTERS, k_size, activation=tf.nn.relu)\n",
    "        _x = tf.reduce_max(_x, 1)\n",
    "        _x = tf.reshape(_x, (batch_size, FILTERS))\n",
    "        feat_map.append(_x)\n",
    "    x = tf.concat(feat_map, -1)\n",
    "    \n",
    "    x = tf.layers.dropout(x, 0.2, training=is_training)\n",
    "    x = tf.layers.dense(x, FILTERS, tf.nn.relu)\n",
    "    logits = tf.layers.dense(x, N_CLASS)\n",
    "    return logits\n",
    "\n",
    "\n",
    "def model_fn(features, labels, mode, params):\n",
    "    logits = forward(features, mode, params)\n",
    "    \n",
    "    if mode == tf.estimator.ModeKeys.PREDICT:\n",
    "        preds = tf.argmax(logits, -1)\n",
    "        return tf.estimator.EstimatorSpec(mode, predictions=preds)\n",
    "    \n",
    "    if mode == tf.estimator.ModeKeys.TRAIN:\n",
    "        global_step = tf.train.get_global_step()\n",
    "\n",
    "        lr_op = tf.train.exponential_decay(LR['start'],\n",
    "                                           global_step,\n",
    "                                           LR['steps'],\n",
    "                                           LR['end']/LR['start'])\n",
    "\n",
    "        loss_op = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(\n",
    "            logits=logits,\n",
    "            labels=labels))\n",
    "\n",
    "        train_op = tf.train.AdamOptimizer(lr_op).minimize(loss_op,\n",
    "                                                          global_step=global_step)\n",
    "\n",
    "        lth = tf.train.LoggingTensorHook({'lr': lr_op},\n",
    "                                          every_n_iter=100)\n",
    "        \n",
    "        return tf.estimator.EstimatorSpec(mode=mode,\n",
    "                                          loss=loss_op,\n",
    "                                          train_op=train_op,\n",
    "                                          training_hooks=[lth])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def main():\n",
    "    (X_train, y_train), (X_test, y_test) = tf.keras.datasets.imdb.load_data(num_words=VOCAB_SIZE)\n",
    "    idx2word, word2idx = get_idx2word()\n",
    "\n",
    "    MAX_WORD_LEN = 0\n",
    "    all_words = ''\n",
    "    X_train_chars = []\n",
    "    for x in tqdm(X_train, total=len(X_train), ncols=70):\n",
    "        temp = []\n",
    "        for idx in x:\n",
    "            word = idx2word[idx]\n",
    "            all_words += (word+' ')\n",
    "            if len(word) > MAX_WORD_LEN:\n",
    "                MAX_WORD_LEN = len(word)\n",
    "            temp.append(list(word))\n",
    "        X_train_chars.append(temp)\n",
    "\n",
    "    X_test_chars = [[list(idx2word[w]) for w in x] for x in tqdm(X_test,\n",
    "                                                                 total=len(X_test),\n",
    "                                                                 ncols=70)]\n",
    "\n",
    "    char2idx = {k: i+2 for i, (k, v) in enumerate(Counter(list(all_words)).most_common())}\n",
    "    char2idx['<pad>'] = 0\n",
    "    char2idx['<unk>'] = 1\n",
    "\n",
    "    X_train_feat = make_feats(X_train_chars, char2idx, word2idx, MAX_WORD_LEN)\n",
    "    X_test_feat = make_feats(X_test_chars, char2idx, word2idx, MAX_WORD_LEN)\n",
    "    \n",
    "    \n",
    "    estimator = tf.estimator.Estimator(model_fn, params={'char_vocab_size': len(char2idx),\n",
    "                                                         'max_word_len': MAX_WORD_LEN,})\n",
    "\n",
    "    for _ in range(N_EPOCH):\n",
    "        estimator.train(tf.estimator.inputs.numpy_input_fn(\n",
    "            x = X_train_feat, y = y_train,\n",
    "            batch_size = BATCH_SIZE,\n",
    "            shuffle = True))\n",
    "        y_pred = np.fromiter(estimator.predict(tf.estimator.inputs.numpy_input_fn(\n",
    "            x = X_test_feat,\n",
    "            batch_size = BATCH_SIZE,\n",
    "            shuffle = False)), np.int32)\n",
    "        print(\"\\nValidation Accuracy: %.4f\\n\" % (y_pred==y_test).mean())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████| 25000/25000 [00:12<00:00, 2042.53it/s]\n",
      "100%|█████████████████████████| 25000/25000 [00:07<00:00, 3387.85it/s]\n",
      "100%|█████████████████████████| 25000/25000 [00:12<00:00, 1946.03it/s]\n",
      "100%|█████████████████████████| 25000/25000 [00:12<00:00, 1953.61it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using default config.\n",
      "WARNING:tensorflow:Using temporary folder as model directory: /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmpso3fv26t\n",
      "INFO:tensorflow:Using config: {'_model_dir': '/var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmpso3fv26t', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': 600, '_session_config': None, '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_train_distribute': None, '_service': None, '_cluster_spec': <tensorflow.python.training.server_lib.ClusterSpec object at 0x12517acf8>, '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Calling model_fn.\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "INFO:tensorflow:Create CheckpointSaverHook.\n",
      "INFO:tensorflow:Graph was finalized.\n",
      "INFO:tensorflow:Running local_init_op.\n",
      "INFO:tensorflow:Done running local_init_op.\n",
      "INFO:tensorflow:Saving checkpoints for 1 into /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmpso3fv26t/model.ckpt.\n",
      "INFO:tensorflow:loss = 0.68802404, step = 1\n",
      "INFO:tensorflow:lr = 0.005\n",
      "INFO:tensorflow:global_step/sec: 1.2996\n",
      "INFO:tensorflow:loss = 0.44669238, step = 101 (76.948 sec)\n",
      "INFO:tensorflow:lr = 0.004288479 (76.948 sec)\n",
      "INFO:tensorflow:global_step/sec: 1.32699\n",
      "INFO:tensorflow:loss = 0.35037428, step = 201 (75.359 sec)\n",
      "INFO:tensorflow:lr = 0.0036782112 (75.359 sec)\n",
      "INFO:tensorflow:global_step/sec: 1.29582\n",
      "INFO:tensorflow:loss = 0.24762842, step = 301 (77.171 sec)\n",
      "INFO:tensorflow:lr = 0.0031547868 (77.171 sec)\n",
      "INFO:tensorflow:global_step/sec: 1.31759\n",
      "INFO:tensorflow:loss = 0.39225727, step = 401 (75.896 sec)\n",
      "INFO:tensorflow:lr = 0.0027058476 (75.896 sec)\n",
      "INFO:tensorflow:global_step/sec: 1.34434\n",
      "INFO:tensorflow:loss = 0.40967354, step = 501 (74.386 sec)\n",
      "INFO:tensorflow:lr = 0.0023207942 (74.386 sec)\n",
      "INFO:tensorflow:global_step/sec: 1.33406\n",
      "INFO:tensorflow:loss = 0.19487444, step = 601 (74.959 sec)\n",
      "INFO:tensorflow:lr = 0.0019905358 (74.959 sec)\n",
      "INFO:tensorflow:global_step/sec: 1.33156\n",
      "INFO:tensorflow:loss = 0.14160585, step = 701 (75.100 sec)\n",
      "INFO:tensorflow:lr = 0.0017072745 (75.100 sec)\n",
      "INFO:tensorflow:Saving checkpoints for 782 into /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmpso3fv26t/model.ckpt.\n",
      "INFO:tensorflow:Loss for final step: 0.27130345.\n",
      "INFO:tensorflow:Calling model_fn.\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "INFO:tensorflow:Graph was finalized.\n",
      "INFO:tensorflow:Restoring parameters from /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmpso3fv26t/model.ckpt-782\n",
      "INFO:tensorflow:Running local_init_op.\n",
      "INFO:tensorflow:Done running local_init_op.\n",
      "\n",
      "Validation Accuracy: 0.8915\n",
      "\n",
      "INFO:tensorflow:Calling model_fn.\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "INFO:tensorflow:Create CheckpointSaverHook.\n",
      "INFO:tensorflow:Graph was finalized.\n",
      "INFO:tensorflow:Restoring parameters from /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmpso3fv26t/model.ckpt-782\n",
      "INFO:tensorflow:Running local_init_op.\n",
      "INFO:tensorflow:Done running local_init_op.\n",
      "INFO:tensorflow:Saving checkpoints for 783 into /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmpso3fv26t/model.ckpt.\n",
      "INFO:tensorflow:loss = 0.15324135, step = 783\n",
      "INFO:tensorflow:lr = 0.0015053472\n",
      "INFO:tensorflow:global_step/sec: 1.30027\n",
      "INFO:tensorflow:loss = 0.2404376, step = 883 (76.908 sec)\n",
      "INFO:tensorflow:lr = 0.00129113 (76.908 sec)\n",
      "INFO:tensorflow:global_step/sec: 1.36424\n",
      "INFO:tensorflow:loss = 0.18667452, step = 983 (73.301 sec)\n",
      "INFO:tensorflow:lr = 0.001107397 (73.301 sec)\n",
      "INFO:tensorflow:global_step/sec: 1.3514\n",
      "INFO:tensorflow:loss = 0.28593192, step = 1083 (73.997 sec)\n",
      "INFO:tensorflow:lr = 0.00094980985 (73.997 sec)\n",
      "INFO:tensorflow:global_step/sec: 1.34303\n",
      "INFO:tensorflow:loss = 0.103726365, step = 1183 (74.458 sec)\n",
      "INFO:tensorflow:lr = 0.00081464805 (74.458 sec)\n",
      "INFO:tensorflow:global_step/sec: 1.32555\n",
      "INFO:tensorflow:loss = 0.17544705, step = 1283 (75.440 sec)\n",
      "INFO:tensorflow:lr = 0.0006987203 (75.440 sec)\n",
      "INFO:tensorflow:global_step/sec: 1.31127\n",
      "INFO:tensorflow:loss = 0.23709458, step = 1383 (76.262 sec)\n",
      "INFO:tensorflow:lr = 0.00059928955 (76.262 sec)\n",
      "INFO:tensorflow:global_step/sec: 1.29828\n",
      "INFO:tensorflow:loss = 0.14588615, step = 1483 (77.025 sec)\n",
      "INFO:tensorflow:lr = 0.00051400816 (77.026 sec)\n",
      "INFO:tensorflow:Saving checkpoints for 1564 into /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmpso3fv26t/model.ckpt.\n",
      "INFO:tensorflow:Loss for final step: 0.24106644.\n",
      "INFO:tensorflow:Calling model_fn.\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "INFO:tensorflow:Graph was finalized.\n",
      "INFO:tensorflow:Restoring parameters from /var/folders/sx/fv0r97j96fz8njp14dt5g7940000gn/T/tmpso3fv26t/model.ckpt-1564\n",
      "INFO:tensorflow:Running local_init_op.\n",
      "INFO:tensorflow:Done running local_init_op.\n",
      "\n",
      "Validation Accuracy: 0.8980\n",
      "\n"
     ]
    }
   ],
   "source": [
    "main()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
